package kb;
import java.io.File;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;

import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

import dom.Fact;
import dom.HyperAlertType;
import dom.Implication;
import dom.Predicate;


public class KnowledgeBase {
	private static KnowledgeBase instance = null;
	private final static boolean debug = false;
	private final String KNOWLEDGE_BASE_FILE = "knowledgebase/knowlegdeBase.xml";
	
	private Map<Integer,HyperAlertType> hyperAlertTypeHash = new HashMap<Integer,HyperAlertType>();
	private NodeList predicateNodeList;
	private Map<Integer,Predicate> predicateHash = new HashMap<Integer,Predicate>();
	private NodeList implicationNodeList;
	private List<Implication> implications; //should be a list of implications not predicates
	private NodeList hyperAlertTypeNodeList;
	private List<HyperAlertType> hyperAlertTypes;

	protected KnowledgeBase(){
		load();
	}


	public static KnowledgeBase getInstance(){
		if(instance == null){
			instance = new KnowledgeBase();
		}
		return instance;
	}
	/**
	 * Loads data from XML files consisting of the snort knowledge base
	 */

	public void load(){
		
		
		File knowledgeBaseXML;
		Document document = null;
		try{
			knowledgeBaseXML = new File(KNOWLEDGE_BASE_FILE);

			DocumentBuilderFactory dbFactory = DocumentBuilderFactory.newInstance();
			DocumentBuilder dBuilder = dbFactory.newDocumentBuilder();
			document = dBuilder.parse(knowledgeBaseXML);
		}
		catch(Exception ex){
			ex.printStackTrace();
			System.out.print("\n");
		}
		document.getDocumentElement().normalize();
		
		// Load the predicates
		predicateNodeList = document.getElementsByTagName("predicate");
		predicateHash = processPredicateNodeList(predicateNodeList);

		// Load the implications
		implicationNodeList = document.getElementsByTagName("implication");
		implications = processImplicationNodeList(implicationNodeList);

		// Load the hyper alert type
		hyperAlertTypeNodeList = document.getElementsByTagName("hyperAlertType");
		hyperAlertTypes = processHyperAlertTypeNodeList(hyperAlertTypeNodeList);
		
		// Info
		print("Loaded " +predicateHash.size() + " predicates");
		print("Loaded " +implications.size() + " implications");
		print("Loaded " +hyperAlertTypes.size() + " hyper alert types");
	}

	/**
	 * Handles creating hyper alert types
	 * @param hyperAlertNodeList, predicates
	 * @return
	 */
	private List<HyperAlertType> processHyperAlertTypeNodeList(NodeList hyperAlertNodeList){

		List<HyperAlertType> hyperAlertTypes = new ArrayList<HyperAlertType>();

		for (int count = 0; count < hyperAlertNodeList.getLength(); count++) {
			
			Node hyperAlertNode = hyperAlertNodeList.item(count);
			HyperAlertType hat = new HyperAlertType();
			if (hyperAlertNode.getNodeType() == Node.ELEMENT_NODE){
				if(hyperAlertNode.getParentNode() != null){
					// Make sure that we only load the predicate nodes under tag <predicates>
					if(hyperAlertNode.getParentNode().getNodeName().equals("hyperAlertTypes")) {
						Element hyperAlertTypeNode = (Element) hyperAlertNode;
						int id = Integer.parseInt(hyperAlertTypeNode.getAttribute("id"));
						String name = hyperAlertTypeNode.getAttribute("name");
						hat.setId(id);
						hat.setName(name);
						
						// Get the facts under the HyperAlertType nodes
						List<Fact> facts = new LinkedList<Fact>();
						
						NodeList factsNodeList = hyperAlertTypeNode.getElementsByTagName("facts");
						Element fact = (Element) factsNodeList.item(0);
						NodeList factNodeList = fact.getElementsByTagName("fact");
						
						for(int counter2=0; counter2< factNodeList.getLength(); counter2++){
							Node factNode = (Element) factNodeList.item(counter2);
							if(factNode.getNodeType() == Node.ELEMENT_NODE){
								if(factNode != null){
									Element factElement = (Element) factNode;
									facts.add(new Fact(factElement.getAttribute("name")));
								}
							}
						}

						// Get the prerequisites under the HyperAlertType nodes
						NodeList predicateNodeList;
						
						NodeList prerequisiteNodeList = hyperAlertTypeNode.getElementsByTagName("prerequisites");
						Element pre = (Element) prerequisiteNodeList.item(0);
						predicateNodeList = pre.getElementsByTagName("predicate");
						List<Predicate> prerequisites = processPredicateChildren(predicateNodeList);

						// Get the consequences under the HyperAlertType nodes
						NodeList consequencesNodeList = hyperAlertTypeNode.getElementsByTagName("consequences");
						Element con = (Element) consequencesNodeList.item(0);
						predicateNodeList = con.getElementsByTagName("predicate");
						List<Predicate> consequences = processPredicateChildren(predicateNodeList);
						
						HyperAlertType h = new HyperAlertType(name,prerequisites,consequences,facts);

						hyperAlertTypeHash.put(new Integer(id), h);
						hyperAlertTypes.add(h);
						print(h.toString());
					}
				}
			}
		}
		return hyperAlertTypes;
	}

	private List<Predicate> processPredicateChildren(NodeList predicateNodes){
		List<Predicate> predicates = new LinkedList<Predicate>();
		if(predicateNodes == null){
			return null;
		}
		for(int test=0; test< predicateNodes.getLength(); test++){
			Node predicateNode = predicateNodes.item(test);
			if(predicateNode.getNodeType() == Node.ELEMENT_NODE){
				String factName = "sourceIP";
				if(predicateNode.hasAttributes()){
					factName = ((Element)predicateNode).getAttribute("factName");
				}
				if(predicateNode.hasChildNodes()){		
					Node node = predicateNode.getFirstChild();
					if(!node.getNodeValue().trim().isEmpty()){
						int predicateId = Integer.parseInt(node.getNodeValue());
						Predicate p = predicateHash.get(new Integer(predicateId)).clone();
						p.addFact(factName);
						if (p != null){
							predicates.add(p);
						}
					}
				}
			}
		}
		return predicates;
	}

	public HyperAlertType getHyperAlertType(int sigId){
		KnowledgeBase.getInstance();// make sure the knowledgebase is loaded
		return hyperAlertTypeHash.get(new Integer(sigId));
	}
	
	/**
	 * Handles the predicate node list under the tge predicates
	 * @param predicateNodeList
	 * @return
	 */
	private Map<Integer,Predicate> processPredicateNodeList(NodeList predicateNodeList){

		Map<Integer,Predicate> predicateHash = new HashMap<Integer,Predicate>();

		for (int count = 0; count < predicateNodeList.getLength(); count++) {
			Node predicateNode = predicateNodeList.item(count);
			//print("Current Element :" + predicateNode.getNodeName());
			if (predicateNode.getNodeType() == Node.ELEMENT_NODE){
				if(predicateNode.getParentNode() != null){
					// Make sure that we only load the predicate nodes under tag <predicates>
					if(predicateNode.getParentNode().getNodeName().equals("predicates")) {
						Element eElement = (Element) predicateNode;
						int id = Integer.parseInt(eElement.getAttribute("id"));
						String name = eElement.getAttribute("name");
						//int factTypeId = Integer.parseInt(eElement.getAttribute("factType"));
						//print("Id:\t" + id);
						//print("Name:\t" + name);
						predicateHash.put(id, new Predicate(id,name));
					}
				}
			}
		}
		return predicateHash;
	}

	private List<Implication> processImplicationNodeList(NodeList implicationNodeList){

		List<Implication> implications = new ArrayList<Implication>();

		for (int count = 0; count < implicationNodeList.getLength(); count++) {
			Node implicationNode = implicationNodeList.item(count);
			//print("Current Element :" + predicateNode.getNodeName());
			if (implicationNode.getNodeType() == Node.ELEMENT_NODE){
				if(implicationNode.getParentNode() != null){
					// Make sure that we only load the predicate nodes under tag <implications>
					if(implicationNode.getParentNode().getNodeName().equals("implications")) {
						//int implicationId = Integer.parseInt(((Element)implicationNode).getAttribute("id"));
						
						Predicate implying = null;
						
						NodeList implyingNodeList = ((Element) implicationNode).getElementsByTagName("predicate");
						Element predicateNode = (Element)implyingNodeList.item(0);
						
						if(predicateNode != null && predicateNode.hasChildNodes()){		
							implying = predicateHash.get(Integer.parseInt(predicateNode.getFirstChild().getNodeValue().trim()));
						}
						
						NodeList impliedNodeList = ((Element) implicationNode).getElementsByTagName("implied");
						Set<Predicate> impliedSet = new HashSet<Predicate>();
						for(int i = 0; i < impliedNodeList.getLength(); i++){
							if(impliedNodeList.item(0) != null && impliedNodeList.item(0).hasChildNodes())
							impliedSet.add(predicateHash.get(Integer.parseInt(impliedNodeList.item(0).getFirstChild().getNodeValue().trim())));
						}
						if(implying != null && impliedSet.size() > 0){
							implications.add(new Implication(implying,impliedSet));
						}
					}
				}
			}
		}
		return implications;
	}

	public Predicate getPredicate(int id){
		KnowledgeBase.getInstance();// make sure the knowledgebase is loaded
		return predicateHash.get(new Integer(id));
	}
	
	public List<Implication> getImplications(){
		KnowledgeBase.getInstance();// make sure the knowledgebase is loaded
		return implications;
	}
	
	public List<HyperAlertType> getHyperAlertTypes(){
		return hyperAlertTypes;
	}

	private void print(String value){
		if(debug){
			System.out.println(value);
		}
	}
}
